package cn.t.dynamic.switcher.router;

import cn.t.dynamic.switcher.type.DataSourceType;

import javax.sql.DataSource;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 数据源上下文辅助类(设置key)<br/>
 * 注意!本类的数据源切换不能在容器初始化过程中使用,必须等待容器初始化完成后才能使用
 */
public final class DynamicDataSourceContextHolder {

    /**
     * 实际的数据源
     */
    private static final Map<String, DataSource> DS_MAP = new ConcurrentHashMap<>();

    /**
     * 数据源类型-数据源key列表映射
     */
    private static final Map<DataSourceType, List<String>> KEY_MAP = new ConcurrentHashMap<>();

    /**
     * 默认主数据源的key
     */
    private static final ThreadLocal<String> CONTEXT = ThreadLocal.withInitial(() -> null);

    /**
     * 添加数据源
     * @param beanName 数据源BeanName
     * @param type 数据源类型
     * @param dataSource 数据源
     */
    public synchronized static void addMap(String beanName, DataSourceType type, DataSource dataSource) {
        var list = KEY_MAP.get(type);
        if (list == null) {
            list = new ArrayList<String>();
        }
        list.add(beanName);
        KEY_MAP.put(type, list);
        // 实际数据源的缓存
        DS_MAP.put(beanName, dataSource);
    }

    /**
     * 获取数据源
     * @param key 数据源key
     * @return 数据源
     */
    protected synchronized static DataSource getDataSource(String key) {
        return DS_MAP.get(key);
    }

    /**
     * 将第一个从数据源转换为主数据源
     * @return 主数据源的key
     */
    public synchronized static String convertSlaveToMaster() {
        var sList = KEY_MAP.get(DataSourceType.SLAVE);
        var mList = new ArrayList<String>(1);
        String key = sList.remove(0);
        mList.add(key);
        KEY_MAP.put(DataSourceType.MASTER, mList);
        return key;
    }

    /**
     * 当前数据源的key
     * @return 当前数据源的key
     */
    public static String current() {
        return CONTEXT.get();
    }

    /**
     * 切换为从数据源(随机)
     */
    public static void slave() {
        var list = KEY_MAP.get(DataSourceType.SLAVE);
        if (list == null || list.isEmpty()) {
            CONTEXT.set(null);
        } else {
            // 简单随机
            CONTEXT.set(list.get(new Random().nextInt(128) % list.size()));
        }
    }
	
    /**
     * 切换为从数据源(指定)
     * @param key 指定数据源key
     */
    public static void slave(String key) {
        if (KEY_MAP.get(DataSourceType.SLAVE).contains(key)) {
            CONTEXT.set(key);
        } else {
            throw new IllegalArgumentException("no slave datasource matched!...");
        }
    }

    /**
     * 切换为主数据源
     */
    public static void master() {
        // 一定会有一个主数据源
        CONTEXT.set(KEY_MAP.get(DataSourceType.MASTER).get(0));
    }

    /**
     * 指定数据源
     * @param key 指定数据源key
     */
    public static void specify(String key) {
        if (DS_MAP.containsKey(key)) {
            CONTEXT.set(key);
        } else {
            throw new IllegalArgumentException("key not matched!");
        }
    }

    /**
     * 重置数据源
     */
    public static void reset() {
        // 如果找不到对应的key,AbstractRoutingDataSource中会使用默认的数据源
        // 默认数据源需要在设置AbstractRoutingDataSource时指定
        CONTEXT.remove();
    }
}